#!/usr/bin/env bash
# 文件名: check_go2_env.sh  （Debian + python3.9 实测通过）
set -euo pipefail

RED='\033[31m'; GREEN='\033[32m'; YELLOW='\033[33m'; BLUE='\033[34m'; NC='\033[0m'
log_info()    { echo -e "${GREEN}[INFO]${NC} $*"; }
log_warn()    { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error()   { echo -e "${RED}[ERROR]${NC} $*"; }
log_debug()   { echo -e "${BLUE}[DEBUG]${NC} $*"; }
exit_fail()   { log_error "Go2 环境检查未通过！"; exit 1; }
exit_ok()     { log_info  "Go2 环境已就绪！"; exit 0; }

# -------------- 1. 找到 Python ≥ 3.8 --------------
locate_python() {
    [[ -n "${PYTHON_CMD:-}" ]] && \
        { verify_python_version "$PYTHON_CMD" && return 0; } || true

    for cmd in python3.11 python3.10 python3.9 python3.8 python3; do
        command -v "$cmd" &>/dev/null || continue
        verify_python_version "$cmd" && return 0
    done
    log_error "未找到 Python ≥ 3.8"; return 1
}

verify_python_version() {
    local py="$1"
    local ver_str major
    ver_str=$("$py" --version 2>&1)          # "Python 3.9.16"
    major=${ver_str#Python }; major=${major%.*}
    if awk -v v="$major" 'BEGIN{exit !(v >= 3.8)}'; then
        export PYTHON_CMD="$py"
        log_info "Python 版本: $major (符合 ≥3.8)"
        return 0
    else
        log_error "$py 版本 $major 低于 3.8"; return 1
    fi
}

# -------------- 2. 检查 pip 包（宽松版本检查） --------------
verify_pip_pkg() {
    local pkg_spec=$1
    local pkg_name=${pkg_spec%%[=<>!]*}
    local strict_check=${2:-0}  # 默认不严格检查版本
    
    log_info "检测 pip 包: $pkg_spec"

    "$PYTHON_CMD" - "$pkg_name" "$pkg_spec" "$strict_check" <<'EOF'
import sys, importlib.metadata as im
pkg_name, pkg_spec, strict_check = sys.argv[1], sys.argv[2], int(sys.argv[3])

try:
    installed = im.version(pkg_name)
    
    if "==" in pkg_spec:
        required = pkg_spec.split("==")[1]
        if installed == required:
            print(f"[OK] {pkg_spec} 已安装")
        else:
            if strict_check == 1:
                print(f"[FAIL] 版本不匹配: 需要 {required}, 实际 {installed}")
                sys.exit(1)
            else:
                print(f"[WARN] 版本不同但可用: 需要 {required}, 实际 {installed}")
    else:
        print(f"[OK] {pkg_name} 已安装（版本: {installed}）")
        
except im.PackageNotFoundError:
    # 特殊处理：如果包未找到，但包名是 opencv-python，检查是否有其他 opencv 包
    if pkg_name == "opencv-python":
        # 检查是否有任何 opencv 相关的包
        opencv_variants = ["opencv-python", "opencv-contrib-python", "opencv-python-headless"]
        for variant in opencv_variants:
            try:
                ver = im.version(variant)
                print(f"[OK] 找到替代包 {variant}（版本: {ver}）")
                sys.exit(0)
            except im.PackageNotFoundError:
                continue
        
        # 检查是否能导入 cv2
        try:
            import cv2
            cv2_version = cv2.__version__
            print(f"[OK] OpenCV 可用（通过系统安装，版本: {cv2_version}）")
            sys.exit(0)
        except ImportError:
            print(f"[FAIL] {pkg_name} 未安装且无法导入 OpenCV")
            sys.exit(1)
    else:
        print(f"[FAIL] {pkg_name} 未安装")
        sys.exit(1)
EOF
}

# -------------- 3. 检查系统 OpenCV 安装情况 --------------
check_system_opencv() {
    log_info "检查系统 OpenCV 安装情况..."
    
    # 检查通过 apt 安装的 OpenCV 包
    local apt_packages
    apt_packages=$(dpkg -l | grep -E "(opencv|libopencv)" | awk '{print $2 "\t" $3}' 2>/dev/null || true)
    
    if [[ -n "$apt_packages" ]]; then
        log_info "通过 APT 安装的 OpenCV 相关包:"
        echo "$apt_packages" | while read -r pkg version; do
            echo "  📦 $pkg: $version"
        done
        return 0
    else
        log_warn "未通过 APT 安装 OpenCV 系统包"
        return 1
    fi
}

# -------------- 4. 检查 Python OpenCV 详细信息 --------------
check_python_opencv_details() {
    log_info "检查 Python OpenCV 详细信息..."
    
    if "$PYTHON_CMD" - <<'EOF'
import sys
try:
    import cv2
    print("✅ Python OpenCV 信息:")
    print(f"   版本: {cv2.__version__}")
    
    # 检查模块路径
    import os
    cv2_path = os.path.abspath(cv2.__file__)
    print(f"   模块路径: {cv2_path}")
    
    # 检查可用的功能
    print("   基本功能测试:")
    try:
        # 创建测试图像
        import numpy as np
        test_img = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
        
        # 测试基本操作
        gray = cv2.cvtColor(test_img, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(test_img, (5, 5), 0)
        edges = cv2.Canny(test_img, 50, 150)
        
        print("     ✅ 图像处理功能正常")
        print(f"     ✅ 图像尺寸: {test_img.shape}")
        print(f"     ✅ 灰度图尺寸: {gray.shape}")
        sys.exit(0)  # 成功退出
        
    except Exception as e:
        print(f"     ❌ 图像处理测试失败: {e}")
        sys.exit(1)  # 失败退出
        
except ImportError as e:
    print(f"❌ 无法导入 OpenCV: {e}")
    sys.exit(1)
except Exception as e:
    print(f"❌ 检查 OpenCV 时出错: {e}")
    sys.exit(1)
EOF
    then
        return 0  # Python 代码执行成功
    else
        return 1  # Python 代码执行失败
    fi
}

# -------------- 5. 综合环境评估 --------------
evaluate_environment() {
    local python_opencv_ok=$1
    local cyclonedds_ok=$2
    local numpy_ok=$3
    local opencv_pip_ok=$4
    
    log_info "=== 环境评估结果 ==="
    
    # 核心依赖检查
    if [ $python_opencv_ok -eq 0 ]; then
        log_info "✅ OpenCV 功能: 正常可用"
    else
        log_error "❌ OpenCV 功能: 不可用"
        return 1
    fi
    
    if [ $numpy_ok -eq 0 ]; then
        log_info "✅ NumPy: 已安装"
    else
        log_error "❌ NumPy: 缺失"
        return 1
    fi
    
    if [ $cyclonedds_ok -eq 0 ]; then
        log_info "✅ CycloneDDS: 已安装（版本可能不同但可用）"
    else
        log_warn "⚠️  CycloneDDS: 未安装"
        # CycloneDDS 不是核心依赖，不因此失败
    fi
    
    # 总结
    if [ $python_opencv_ok -eq 0 ] && [ $numpy_ok -eq 0 ]; then
        log_info "🎉 核心依赖 (OpenCV + NumPy) 已就绪，Go2 环境可用！"
        return 0
    else
        log_error "核心依赖不完整，请检查上述错误信息"
        return 1
    fi
}

# -------------- 主流程 --------------
main() {
    locate_python || exit_fail

    log_info "=== 开始环境检查 ==="
    
    # 检查系统 OpenCV
    check_system_opencv || true
    
    echo
    
    # 检查 Python OpenCV 详细信息
    PYTHON_OPENCV_OK=0
    if check_python_opencv_details; then
        PYTHON_OPENCV_OK=0
        log_info "✅ Python OpenCV 检查完成"
    else
        PYTHON_OPENCV_OK=1
        log_error "❌ Python OpenCV 检查失败"
    fi
    
    echo
    
    # 检查必要的 pip 包（宽松检查）
    log_info "检查必要的 pip 包..."
    
    CYCLONEDDS_OK=0
    NUMPY_OK=0
    OPENCV_PIP_OK=0
    
    # CycloneDDS: 版本不同但可用
    if verify_pip_pkg "cyclonedds==0.10.2" 0; then
        CYCLONEDDS_OK=0
    else
        CYCLONEDDS_OK=1
    fi
    
    # NumPy: 必须存在
    if verify_pip_pkg "numpy" 1; then
        NUMPY_OK=0
    else
        NUMPY_OK=1
    fi
    
    # OpenCV: 宽松检查，只要能用就行
    if verify_pip_pkg "opencv-python" 0; then
        OPENCV_PIP_OK=0
    else
        OPENCV_PIP_OK=1
    fi
    
    echo
    
    # 综合评估
    if evaluate_environment $PYTHON_OPENCV_OK $CYCLONEDDS_OK $NUMPY_OK $OPENCV_PIP_OK; then
        exit_ok
    else
        exit_fail
    fi
}

main "$@"